import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏窗口大小
WIDTH = 600
HEIGHT = 400
CELL_SIZE = 20

# 颜色定义
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLACK = (0, 0, 0)

# 初始化窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇")

# 时钟与字体
clock = pygame.time.Clock()
small_font = pygame.font.SysFont(None, 24)
large_font = pygame.font.SysFont(None, 48)

# 状态常量
STATE_RUNNING = "RUNNING"
STATE_PAUSED = "PAUSED"
STATE_GAME_OVER = "GAME_OVER"

# 记录与速度
score = 0
high_score = 0
speed = 10  # 初始速度（FPS）

# 蛇与食物初始化
def reset_game():
    global snake, snake_dir, food, score, speed, state
    snake = [(100, 100), (80, 100), (60, 100)]
    snake_dir = (CELL_SIZE, 0)  # 初始向右移动
    food = spawn_food()
    score = 0
    speed = 10
    state = STATE_RUNNING


def spawn_food():
    """在不与蛇身重叠的位置生成食物。"""
    # 网格所有位置
    cols = WIDTH // CELL_SIZE
    rows = HEIGHT // CELL_SIZE
    all_positions = [(x * CELL_SIZE, y * CELL_SIZE) for x in range(cols) for y in range(rows)]
    snake_set = set(snake) if 'snake' in globals() else set()
    available = [pos for pos in all_positions if pos not in snake_set]
    if not available:
        # 没有可用位置（罕见：填满了）——保持原食物或返回None
        return None
    return random.choice(available)


def draw_snake():
    for segment in snake:
        pygame.draw.rect(screen, GREEN, pygame.Rect(segment[0], segment[1], CELL_SIZE, CELL_SIZE))


def draw_food():
    if food is not None:
        pygame.draw.rect(screen, RED, pygame.Rect(food[0], food[1], CELL_SIZE, CELL_SIZE))


def draw_text(text, pos, color=WHITE, center=False, font=None):
    font = font or small_font
    surface = font.render(text, True, color)
    rect = surface.get_rect()
    if center:
        rect.center = pos
    else:
        rect.topleft = pos
    screen.blit(surface, rect)


# 初始化游戏
state = STATE_RUNNING
snake = [(100, 100), (80, 100), (60, 100)]
snake_dir = (CELL_SIZE, 0)
food = spawn_food()

# 游戏循环
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                pygame.quit()
                sys.exit()

            # 暂停/继续
            if event.key == pygame.K_p and state in (STATE_RUNNING, STATE_PAUSED):
                state = STATE_PAUSED if state == STATE_RUNNING else STATE_RUNNING

            # 重开（仅在Game Over）
            if event.key == pygame.K_r and state == STATE_GAME_OVER:
                reset_game()
                continue

            # 方向控制（仅在运行时且不能直接反向）
            if state == STATE_RUNNING:
                if event.key == pygame.K_UP and snake_dir != (0, CELL_SIZE):
                    snake_dir = (0, -CELL_SIZE)
                elif event.key == pygame.K_DOWN and snake_dir != (0, -CELL_SIZE):
                    snake_dir = (0, CELL_SIZE)
                elif event.key == pygame.K_LEFT and snake_dir != (CELL_SIZE, 0):
                    snake_dir = (-CELL_SIZE, 0)
                elif event.key == pygame.K_RIGHT and snake_dir != (-CELL_SIZE, 0):
                    snake_dir = (CELL_SIZE, 0)

    if state == STATE_RUNNING:
        # 移动蛇
        new_head = (snake[0][0] + snake_dir[0], snake[0][1] + snake_dir[1])
        snake.insert(0, new_head)

        # 吃到食物
        if food is not None and snake[0] == food:
            score += 1
            if score > high_score:
                high_score = score
            # 速度随得分提升
            speed = min(60, speed + 1)  # 上限避免过快
            # 刷新新食物（避免生成到蛇身上）
            food = spawn_food()
        else:
            snake.pop()

        # 撞墙或咬到自己
        hit_wall = (snake[0][0] < 0 or snake[0][0] >= WIDTH or
                    snake[0][1] < 0 or snake[0][1] >= HEIGHT)
        hit_self = snake[0] in snake[1:]
        if hit_wall or hit_self:
            state = STATE_GAME_OVER

    # 绘制画面
    screen.fill(BLACK)
    draw_snake()
    draw_food()

    # HUD：分数与最高分
    draw_text(f"Score: {score}  High: {high_score}", (10, 10), WHITE, center=False, font=small_font)

    # 状态提示
    if state == STATE_PAUSED:
        draw_text("Paused (P 继续)", (WIDTH // 2, HEIGHT // 2 - 20), color=WHITE, center=True, font=large_font)
    elif state == STATE_GAME_OVER:
        draw_text("Game Over", (WIDTH // 2, HEIGHT // 2 - 30), color=RED, center=True, font=large_font)
        draw_text("按 R 重开 / ESC 退出", (WIDTH // 2, HEIGHT // 2 + 10), color=WHITE, center=True, font=small_font)

    pygame.display.flip()

    # 控制速度
    if state == STATE_RUNNING:
        clock.tick(speed)
    else:
        clock.tick(10)
